import os, sys
import baostock as bs
from datetime import datetime, timedelta
import pandas as pd
from utils.common import *


def download_all_stock_names(date_string=''):
    ''' 读取数据，并保存到 all_stock_names_{date_string}.csv。
        df，包括 id, code, tradeStatus 和 code_name。如：
                       code   tradeStatus      code_name
            0     sh.000001             1      上证综合指数
            1     sh.000002             1      上证A股指数
        其中，tradeStatus表示交易状态，1：正常交易 0：停牌。
        返回：输出文件CSV的路径。
        注意: 节假日返回为0，即没有数据。
        参考：http://baostock.com/baostock/index.php/Python_API%E6%96%87%E6%A1%A3
    '''
    lg = bs.login()
    if  lg.error_code != '0':
        print("Loging failed. Error code: " + lg.error_msg)
        print("Loging failed. Error message: " + lg.error_code)
        return
    
    try:
        date_string = '2025-01-02' if len(date_string) == 0 else date_string
        d1 = datetime.strptime(date_string, '%Y-%m-%d')
    except Exception:
        date_string = f'{datetime.now():%Y-%m-%d}'
    
    # 获取指定日期的指数、股票数据
    print('date:', date_string)
    stock_rs = bs.query_all_stock(date_string)
    stock_df = stock_rs.get_data()
    bs.logout()
    
    filename = f'all_stock_codes.csv'
    stock_df.to_csv(filename, encoding="utf-8", index=False)
    #print('all stocks saved to: allstocks.csv, size:', len(stock_df))
    return filename